{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import re\n",
    "os.environ['TOKENIZERS_PARALLELISM'] = \"false\"\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '1'\n",
    "import sys\n",
    "sys.path.insert(1, os.path.join(sys.path[0], './utils'))\n",
    "sys.path.insert(1, os.path.join(sys.path[0], './pytorch'))\n",
    "import numpy as np\n",
    "import argparse\n",
    "import time\n",
    "import logging\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torch.utils.data\n",
    " \n",
    "from utilities import (create_folder, get_filename, create_logging, Mixup, \n",
    "    StatisticsContainer, split_phrases, pad_or_truncate)\n",
    "from models import (Cnn14, Cnn14_no_specaug, Cnn14_no_dropout, \n",
    "    Cnn6, Cnn10, ResNet22, ResNet38, ResNet54, Cnn14_emb512, Cnn14_emb128, \n",
    "    Cnn14_emb32, MobileNetV1, MobileNetV2, LeeNet11, LeeNet24, DaiNet19, \n",
    "    Res1dNet31, Res1dNet51, Wavegram_Cnn14, Wavegram_Logmel_Cnn14, \n",
    "    Wavegram_Logmel128_Cnn14, Cnn14_16k, Cnn14_8k, Cnn14_mel32, Cnn14_mel128, \n",
    "    Cnn14_mixup_time_domain, Cnn14_DecisionLevelMax, Cnn14_DecisionLevelAtt, Cnn14Bert)\n",
    "from pytorch_utils import (move_data_to_device, count_parameters, count_flops, \n",
    "    do_mixup)\n",
    "from data_generator import (AudioSetDataset, TrainSampler, BalancedTrainSampler, \n",
    "    AlternateTrainSampler, EvaluateSampler, collate_fn, AudioSetBiModalDataset, get_collate_fn)\n",
    "from evaluate import Evaluator, EvaluatorBiModal\n",
    "import config\n",
    "import librosa\n",
    "import pandas as pd\n",
    "from collections import defaultdict\n",
    "from tqdm import tqdm, trange\n",
    "import pickle\n",
    "import multiprocessing\n",
    "from sklearn.metrics.pairwise import cosine_similarity\n",
    "from rank_metrics import retrieval_as_classification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at prajjwal1/bert-medium were not used when initializing BertModel: ['cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.decoder.bias', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.decoder.weight', 'cls.predictions.bias', 'cls.seq_relationship.weight', 'cls.seq_relationship.bias']\n",
      "- This IS expected if you are initializing BertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    }
   ],
   "source": [
    "model = Cnn14Bert(sample_rate=32000, window_size=1024, \n",
    "        hop_size=160, mel_bins=64, fmin=50, fmax=8000, \n",
    "        bert_model_type=\"prajjwal1/bert-medium\", \n",
    "        max_seq_len=16, shared_dim=1024)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "resume_checkpoint_path = \"/backup/data/Audioset/zelin/workspaces/audioset_tagging/checkpoints_bimodal/main_bimodal/sample_rate=32000,window_size=1024,hop_size=160,mel_bins=64,fmin=50,fmax=8000/data_type=balanced_train/Cnn14Bert/bert_type=prajjwal1/bert-medium/balanced=balanced/max_text_nums=128/max_seq_len=16/batch_size=16/50000_iterations.pth\"\n",
    "checkpoint = torch.load(resume_checkpoint_path)\n",
    "model.load_state_dict(checkpoint['model'], strict=False)\n",
    "model.cuda()\n",
    "model.eval()\n",
    "None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>caption_1</th>\n",
       "      <th>caption_2</th>\n",
       "      <th>caption_3</th>\n",
       "      <th>caption_4</th>\n",
       "      <th>caption_5</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>file_name</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>Santa Motor.wav</th>\n",
       "      <td>A machine whines and squeals while rhythmicall...</td>\n",
       "      <td>A person is using electric clippers to trim bu...</td>\n",
       "      <td>Someone is trimming the bushes with electric c...</td>\n",
       "      <td>The whirring of a pump fills a bladder that tu...</td>\n",
       "      <td>While rhythmically punching or stamping, a mac...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Radio Garble.wav</th>\n",
       "      <td>A radio dispatcher and an officer are communic...</td>\n",
       "      <td>Communication with a walkie-talkie with a lot ...</td>\n",
       "      <td>A discussion with a walkie-talkie with a consi...</td>\n",
       "      <td>People talking through a walkie-talkie with ba...</td>\n",
       "      <td>The walkie-talkie the people are talking throu...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Radio Fuzz for Old Radio Broadcast FF233.wav</th>\n",
       "      <td>A radio tuner has been positioned in between r...</td>\n",
       "      <td>A transistor radio is being played on a statio...</td>\n",
       "      <td>A transistor radio is on a station that is not...</td>\n",
       "      <td>Radio static makes a constant hum with a high ...</td>\n",
       "      <td>Static coming from a radio that is in between ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>toy rattle 2.wav</th>\n",
       "      <td>A person winding up a device and then jingling...</td>\n",
       "      <td>A socket wrench that is tightening a bolt.</td>\n",
       "      <td>An object is tightened and then metallic objec...</td>\n",
       "      <td>Before keys are jangled on their chain, someon...</td>\n",
       "      <td>Someone is spinning around a lock with a dial.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Blade Big.wav</th>\n",
       "      <td>A person is pulling silverware out of the dish...</td>\n",
       "      <td>A person removes a knife from its holder then ...</td>\n",
       "      <td>A person taking a knife out of its holder and ...</td>\n",
       "      <td>Metal sliding together such as swords or knives.</td>\n",
       "      <td>The metallic clang of swords and knives striki...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                                                      caption_1  \\\n",
       "file_name                                                                                         \n",
       "Santa Motor.wav                               A machine whines and squeals while rhythmicall...   \n",
       "Radio Garble.wav                              A radio dispatcher and an officer are communic...   \n",
       "Radio Fuzz for Old Radio Broadcast FF233.wav  A radio tuner has been positioned in between r...   \n",
       "toy rattle 2.wav                              A person winding up a device and then jingling...   \n",
       "Blade Big.wav                                 A person is pulling silverware out of the dish...   \n",
       "\n",
       "                                                                                      caption_2  \\\n",
       "file_name                                                                                         \n",
       "Santa Motor.wav                               A person is using electric clippers to trim bu...   \n",
       "Radio Garble.wav                              Communication with a walkie-talkie with a lot ...   \n",
       "Radio Fuzz for Old Radio Broadcast FF233.wav  A transistor radio is being played on a statio...   \n",
       "toy rattle 2.wav                                     A socket wrench that is tightening a bolt.   \n",
       "Blade Big.wav                                 A person removes a knife from its holder then ...   \n",
       "\n",
       "                                                                                      caption_3  \\\n",
       "file_name                                                                                         \n",
       "Santa Motor.wav                               Someone is trimming the bushes with electric c...   \n",
       "Radio Garble.wav                              A discussion with a walkie-talkie with a consi...   \n",
       "Radio Fuzz for Old Radio Broadcast FF233.wav  A transistor radio is on a station that is not...   \n",
       "toy rattle 2.wav                              An object is tightened and then metallic objec...   \n",
       "Blade Big.wav                                 A person taking a knife out of its holder and ...   \n",
       "\n",
       "                                                                                      caption_4  \\\n",
       "file_name                                                                                         \n",
       "Santa Motor.wav                               The whirring of a pump fills a bladder that tu...   \n",
       "Radio Garble.wav                              People talking through a walkie-talkie with ba...   \n",
       "Radio Fuzz for Old Radio Broadcast FF233.wav  Radio static makes a constant hum with a high ...   \n",
       "toy rattle 2.wav                              Before keys are jangled on their chain, someon...   \n",
       "Blade Big.wav                                  Metal sliding together such as swords or knives.   \n",
       "\n",
       "                                                                                      caption_5  \n",
       "file_name                                                                                        \n",
       "Santa Motor.wav                               While rhythmically punching or stamping, a mac...  \n",
       "Radio Garble.wav                              The walkie-talkie the people are talking throu...  \n",
       "Radio Fuzz for Old Radio Broadcast FF233.wav  Static coming from a radio that is in between ...  \n",
       "toy rattle 2.wav                                 Someone is spinning around a lock with a dial.  \n",
       "Blade Big.wav                                 The metallic clang of swords and knives striki...  "
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "val_df = pd.read_csv(\"/home/zhiling/py3_workspace/aser_audioset/clotho/eval.csv\", index_col=0)\n",
    "val_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['although room was initially serene',\n",
       " 'people talk',\n",
       " 'laugh with loud person near end']"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cap = \"Although the room was initially serene, people talk and laugh with a loud person near the end.\"\n",
    "pos_phrases = split_phrases(cap)\n",
    "pos_phrases"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['squeals', 'machine whines', 'rhythmically punching or stamping']"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cap = \"A machine whines and squeals while rhythmically punching or stamping.\"\n",
    "neg_phrases = split_phrases(cap)\n",
    "neg_phrases"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_wav(full_path):\n",
    "    (waveform, _) = librosa.core.load(full_path, sr=config.sample_rate, mono=True)\n",
    "    waveform = pad_or_truncate(waveform, config.clip_samples)\n",
    "    return waveform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "wav_dir = \"/home/zhiling/py3_workspace/aser_audioset/zelin_dev/clotho/eval/\"\n",
    "waveform = load_wav(wav_dir+\"young artists.wav\")\n",
    "waveform = waveform[None, :]    # (1, audio_length)\n",
    "waveform = move_data_to_device(waveform, 'cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([1, 1024]), tensor(1., device='cuda:0', grad_fn=<NormBackward0>))"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "audio_emb = model(waveform=waveform)['audio_emb']\n",
    "audio_emb = torch.nn.functional.normalize(audio_emb)\n",
    "audio_emb.shape, audio_emb.norm()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([6, 1024])"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "phrase_embs = model(texts=pos_phrases+neg_phrases)['text_emb']\n",
    "phrase_embs = torch.nn.functional.normalize(phrase_embs)\n",
    "phrase_embs.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.1295, -0.1026, -0.2038, -0.2578, -0.1751, -0.1915]],\n",
       "       device='cuda:0', grad_fn=<MmBackward>)"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 甚至没有办法区分两个相差很大的句子里的phrase?\n",
    "audio_emb @ phrase_embs.T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1.0000, 0.4818, 0.4666, 0.1125, 0.3808, 0.3518],\n",
       "        [0.4818, 1.0000, 0.6172, 0.1432, 0.3657, 0.2890],\n",
       "        [0.4666, 0.6172, 1.0000, 0.2990, 0.3406, 0.3440],\n",
       "        [0.1125, 0.1432, 0.2990, 1.0000, 0.2024, 0.3336],\n",
       "        [0.3808, 0.3657, 0.3406, 0.2024, 1.0000, 0.4316],\n",
       "        [0.3518, 0.2890, 0.3440, 0.3336, 0.4316, 1.0000]], device='cuda:0',\n",
       "       grad_fn=<MmBackward>)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 单模态的还算合理\n",
    "phrase_embs @ phrase_embs.T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "waveform2 = load_wav(wav_dir+\"Family Reunion Side A Original.wav\")\n",
    "waveform2 = waveform2[None, :]    # (1, audio_length)\n",
    "waveform2 = move_data_to_device(waveform2, 'cuda')\n",
    "audio_emb2 = model(waveform=waveform2)['audio_emb']\n",
    "audio_emb2 = torch.nn.functional.normalize(audio_emb2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.2220]], device='cuda:0', grad_fn=<MmBackward>)"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "audio_emb @ audio_emb2.T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "waveform3 = load_wav(wav_dir+\"Santa Motor.wav\")\n",
    "waveform3 = waveform3[None, :]    # (1, audio_length)\n",
    "waveform3 = move_data_to_device(waveform3, 'cuda')\n",
    "audio_emb3 = model(waveform=waveform3)['audio_emb']\n",
    "audio_emb3 = torch.nn.functional.normalize(audio_emb3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "waveform3 = load_wav(wav_dir+\"Santa Motor.wav\")\n",
    "waveform3 = waveform3[None, :]    # (1, audio_length)\n",
    "waveform3 = move_data_to_device(waveform3, 'cuda')\n",
    "audio_emb3 = model(waveform=waveform3)['audio_emb']\n",
    "audio_emb3 = torch.nn.functional.normalize(audio_emb3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.2853]], device='cuda:0', grad_fn=<MmBackward>)"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "audio_emb @ audio_emb3.T"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "音频的相似度好像有点问题，更接近的audio反而相似度低，不过两条音频的音质之类的差距确实很大"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "74d07dc07065639914d79fc95a011a49c7b91ce6bf0d82680fb06e81dd1e87b0"
  },
  "kernelspec": {
   "display_name": "Python 3.6.10 64-bit ('py36': conda)",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
